{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Molecular dynamics \n",
    "\n",
    "* *simple visualization* of live MD simulation visualisation \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run if you need extra packages: numba and scipy\n",
    "!pip install scipy numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "\n",
    "import numpy as np\n",
    "from MD_utils import lennardjones, simple_molecule_vis\n",
    "from scipy.constants import N_A\n",
    "epsilon = 0.9959\n",
    "sigma = 0.3405\n",
    "k_B = 0.008311\n",
    "M = 39.948"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### initial condition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/marcin/anaconda3/lib/python3.7/site-packages/traittypes/traittypes.py:101: UserWarning: Given trait value dtype \"float64\" does not match required type \"float32\". A coerced copy has been created.\n",
      "  np.dtype(self.dtype).name))\n",
      "/Users/marcin/anaconda3/lib/python3.7/site-packages/traittypes/traittypes.py:101: UserWarning: Given trait value dtype \"int64\" does not match required type \"uint32\". A coerced copy has been created.\n",
      "  np.dtype(self.dtype).name))\n"
     ]
    }
   ],
   "source": [
    "plt = simple_molecule_vis(bs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e8b6b1c18a4478ab9489d2698d02e02",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot.display()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no plot object\n",
      "512\n"
     ]
    }
   ],
   "source": [
    "import numpy  as np\n",
    "# bs jest rozmiarem kostki w której znajdują się cząstki.  \n",
    "nx = 8\n",
    "dx = sigma *1.2\n",
    "bs = dx*nx \n",
    "box = np.array([bs,bs,bs])\n",
    "try:\n",
    "    plt.update_box(bs=bs)\n",
    "    print(\"no plot object\")\n",
    "except:\n",
    "    pass\n",
    "\n",
    "Nparticles = nx**3\n",
    "print(Nparticles)\n",
    "\n",
    "U = np.zeros((3,Nparticles))\n",
    "l = 0\n",
    "for i in range(nx):\n",
    "    for j in range(nx):\n",
    "        for k in range(nx):\n",
    "            U[:,l] = (dx*i-dx*(nx/2-0.50),dx*j-dx*(nx/2-0.5),dx*k-dx*(nx/2-0.5))\n",
    "            l+=1\n",
    "V = np.zeros_like(U)\n",
    "plt.pkts.positions = U.T.copy()\n",
    "plt.pkts.point_size = .3\n",
    "plt.pkts.colors = 0xff0000*np.ones(Nparticles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.43 ms, sys: 168 µs, total: 4.6 ms\n",
      "Wall time: 4.55 ms\n"
     ]
    }
   ],
   "source": [
    "box = np.array([bs,bs,bs])\n",
    "%time Epot, F, Vir = lennardjones(U, box,sigma = 0.3405, epsilon=0.9959)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = 0.1*(np.random.randn(3,Nparticles)-0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1510 -2142.504624736757 0.009524307305879267 -2142.4951004294508 T= 0.0014980236926538445 P= -709.3923955547551\n",
      "Vir: -74331.50788855158\n",
      "bs: 3.2688\n",
      "CPU times: user 6.5 s, sys: 260 ms, total: 6.76 s\n",
      "Wall time: 6.68 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Start simulation\n",
    "\n",
    "import time \n",
    "n_steps = 1520\n",
    "dt = 0.005\n",
    "N = Nparticles\n",
    "\n",
    "box = np.array([bs,bs,bs])\n",
    "plt.update_box(bs=bs)\n",
    "\n",
    "(epot,F,Vir) = lennardjones(U,box,sigma = 0.3405, epsilon=0.9959)\n",
    "traj = []\n",
    "for i in range(n_steps):\n",
    "    U += V * dt + 0.5 * F/M * dt * dt\n",
    "    F0 = F[:]\n",
    "\n",
    "    (epot, F, Vir) = lennardjones(U, box,sigma = 0.3405, epsilon=0.9959)\n",
    "    \n",
    "    V += 0.5 * (F + F0)/M * dt\n",
    "    U -= bs*np.rint(U/bs)\n",
    "    \n",
    "    traj.append(U[:,233].copy())\n",
    "    \n",
    "    if i%10==0:\n",
    "        T = M*np.sum(V**2)/(k_B*(3*N-6))  \n",
    "        P = 1/bs**3*(  N*k_B*T + 1/3.*Vir )\n",
    "        Ek = 0.5*M*np.sum(V**2)\n",
    "        plt.pkts.positions = U.T.copy()\n",
    "        plt.pkts.colors = (np.sum((V**2),axis=0)/10.0* 0xffffff).astype(np.int64)\n",
    "        plt.update_box(bs=bs)\n",
    "        clear_output(wait=True)\n",
    "     \n",
    "        print(i,epot,Ek,epot+Ek,\"T=\",T,\"P=\",P)\n",
    "        print(\"Vir:\",Vir)\n",
    "        print(\"bs:\",bs)\n",
    "    \n",
    "traj = np.array(traj)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The trajectory of a single atom "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import k3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_traj = k3d.points(traj.copy(), color=0xffff00, point_size=.05)\n",
    "plt.plot += plt_traj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  },
  "nbTranslate": {
   "displayLangs": [
    "en",
    "pl"
   ],
   "hotkey": "alt-t",
   "langInMainMenu": true,
   "sourceLang": "pl",
   "targetLang": "en",
   "useGoogleTranslate": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
